# visual_verify_processed_hypersim.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm
import cv2
import torch

# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
# Note: Assumes verify_utils.py is in the same directory or accessible in the python path.
try:
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' from the original verification script is in the same directory.")
    sys.exit(1)


def load_processed_view_data(rgb_path, logger):
    """
    Loads all necessary data for a single view from the preprocessed dataset format.
    """
    try:
        base_path = rgb_path.replace("_rgb.png", "")
        depth_path = f"{base_path}_depth.npy"
        cam_path = f"{base_path}_cam.npz"

        if not all(osp.exists(p) for p in [depth_path, cam_path]):
            logger.warning(f"Missing files for {osp.basename(rgb_path)}. Skipping.")
            return None

        rgb_image = imread_cv2(rgb_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
        if rgb_image is None:
            logger.warning(f"Could not read image {rgb_path}. Skipping.")
            return None
        
        depthmap = np.load(depth_path).astype(np.float32)
        depthmap[~np.isfinite(depthmap)] = 0
        
        depth_h, depth_w = depthmap.shape[:2]
        rgb_h, rgb_w = rgb_image.shape[:2]
        if depth_h != rgb_h or depth_w != rgb_w:
            logger.debug(f"Resizing RGB image from {rgb_image.shape} to match depth {depthmap.shape}")
            rgb_image = cv2.resize(rgb_image, (depth_w, depth_h), interpolation=cv2.INTER_LINEAR)

        cam_file = np.load(cam_path)
        intrinsics = cam_file["intrinsics"].astype(np.float32)
        camera_pose = cam_file["pose"].astype(np.float32)
        
        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of a PREPROCESSED Hypersim dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='E:/data/3D/infinigen_processed',
        help="Path to the root directory of the preprocessed Hypersim dataset (e.g., /path/to/resized_hypersim/train)."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_processed.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='1a1f2004',
        help="Optional: Specify a single scene name (e.g., '22ca3b8b') to visualize."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )
    
    args = parser.parse_args()
    
    logger = setup_logging(args.log_file)

    if not osp.isdir(args.data_root):
        logger.error(f"Provided data root '{args.data_root}' does not exist.")
        sys.exit(1)
    
    if args.scene:
        scene_dirs = [osp.join(args.data_root, args.scene)]
        if not osp.isdir(scene_dirs[0]):
            logger.error(f"Specified scene directory '{scene_dirs[0]}' does not exist.")
            sys.exit(1)
    else:
        scene_dirs = sorted([d for d in glob.glob(osp.join(args.data_root, 'ai_*')) if osp.isdir(d)])
    
    logger.info(f"Found {len(scene_dirs)} scenes to process in '{args.data_root}'.")
    
    for scene_dir in scene_dirs:
        subscene_dirs = sorted([d for d in glob.glob(osp.join(scene_dir, 'cam*')) if osp.isdir(d)])
        for subscene_dir in subscene_dirs:
            subscene_key = f"{osp.basename(scene_dir)}/{osp.basename(subscene_dir)}"
            logger.info(f"\n--- Processing subscene: {subscene_key} ---")

            rgb_paths = sorted(glob.glob(osp.join(subscene_dir, "*_rgb.png")))[:120:5]
            if not rgb_paths:
                logger.warning(f"No RGB images found in {subscene_dir}. Skipping.")
                continue

            # --- Prepare data for the visualizer ---
            pc_list, color_list, conf_list = [], [], []
            focals, pps, Rs, ts = [], [], [], []
            
            logger.info("Loading views and generating point clouds...")
            for rgb_path in tqdm(rgb_paths, file=sys.stdout):
                view_data = load_processed_view_data(rgb_path, logger)
                if view_data:
                    # Generate world-space point cloud for this single view
                    world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                        view_data["depth"], view_data["intrinsics"], view_data["pose"]
                    )
                    pc_list.append(world_points)
                    color_list.append(view_data['rgb'])
                    conf_list.append(valid_mask.astype(np.float32))

                    # Aggregate camera parameters for the viewer
                    intrinsics = view_data['intrinsics']
                    pose = view_data['pose']
                    focals.append(intrinsics[0, 0])
                    pps.append((intrinsics[0, 2], intrinsics[1, 2]))
                    Rs.append(pose[:3, :3])
                    ts.append(pose[:3, 3])
        
            if not pc_list:
                logger.warning(f"No valid data could be loaded for {subscene_key}. Skipping visualization.")
                continue

            cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

            # --- Launch the visualizer ---
            try:
                logger.info(f"Launching interactive visualizer for {subscene_key} at http://localhost:{args.port}")
                logger.info("--> Close the browser tab or press Ctrl+C in the console to continue to the next subscene.")
                launch_visualizer(
                    pc_list=pc_list,
                    color_list=color_list,
                    conf_list=conf_list,
                    cam_dict=cam_dict,
                    port=args.port
                )
            except KeyboardInterrupt:
                logger.info(f"\nVisualizer for {subscene_key} closed, proceeding...")
            except Exception as e:
                logger.error(f"Failed to launch visualizer for {subscene_key}: {e}", exc_info=True)

    logger.info("\n--- Visual verification complete. ---")
